import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from model import CNNs
# from model import vig
from model.layers import Flatten, L2Norm, GeM, Convs, NN, NNN

cuda = True
device = torch.device("cuda" if cuda else "cpu")


class Lsm(nn.Module):
    def __init__(self, backbone, coarse_matching_dim):
        super(Lsm, self).__init__()
        # cnn
        self.sat_backbone, features_dim = CNNs.get_backbone(backbone)
        self.grd_backbone, features_dim = CNNs.get_backbone(backbone)

        # coarse_matching
        self.sat_coarse_matching = nn.Sequential(
                L2Norm(),
                GeM(),
                Flatten(),
                nn.Linear(features_dim, coarse_matching_dim),
                L2Norm()
            )
        self.grd_coarse_matching = nn.Sequential(
                L2Norm(),
                GeM(),
                Flatten(),
                nn.Linear(features_dim, coarse_matching_dim),
                L2Norm()
            )

    def forward(self, sat, grd):
        sat_feature = self.sat_backbone(sat)
        grd_feature = self.grd_backbone(grd)
        sat_coding = self.sat_coarse_matching(sat_feature)
        grd_coding = self.grd_coarse_matching(grd_feature)

        return sat_feature, grd_feature, sat_coding, grd_coding


class PatchDSM(nn.Module):
    def __init__(self, coarse_matching_dim, fine_matching_dim):
        super(PatchDSM, self).__init__()

        self.conv = Convs(coarse_matching_dim, fine_matching_dim)
        self.nn = NN()

    def forward(self, grd_coding, sat_coding):
        sat_coding = F.normalize(sat_coding, dim=1)
        grd_coding = F.normalize(grd_coding, dim=1)
        ori, option_ori = self.conv(grd_coding, sat_coding)
        ori = torch.squeeze(ori, (2, 3))
        option_ori = torch.squeeze(option_ori, 2)
        # ori = F.normalize(ori, dim=0)
        # option_ori = F.normalize(option_ori, dim=1)
        ori, option_ori = self.nn(ori, option_ori)
        return ori, option_ori


class PatchDSM2(nn.Module):
    def __init__(self, coarse_matching_dim, fine_matching_dim):
        super(PatchDSM2, self).__init__()

        self.conv = Convs(coarse_matching_dim, fine_matching_dim)
        self.nn = NNN()

    def forward(self, grd_coding, sat_coding):
        sat_coding = F.normalize(sat_coding, dim=1)
        grd_coding = F.normalize(grd_coding, dim=1)
        ori, option_ori = self.conv(grd_coding, sat_coding)
        ori = torch.squeeze(ori, (2, 3))
        option_ori = torch.squeeze(option_ori, 2)
        # ori = F.normalize(ori, dim=0)
        # option_ori = F.normalize(option_ori, dim=1)
        ori, option_ori = self.nn(ori, option_ori)
        return ori, option_ori

if __name__ == '__main__':
    class Test(nn.Module):
        def __init__(self, backbone, coarse_matching_dim, fine_matching_dim):
            super(Test, self).__init__()
            # cnn
            self.sat_backbone, features_dim = CNNs.get_backbone(backbone)
            # coarse_matching
            self.sat_coarse_matching = nn.Sequential(
                L2Norm(),
                GeM(),
                Flatten(),
                nn.Linear(features_dim, coarse_matching_dim),
                L2Norm()
            )
            self.grd_coarse_matching = nn.Sequential(
                L2Norm(),
                GeM(),
                Flatten(),
                nn.Linear(features_dim, coarse_matching_dim),
                L2Norm()
            )

